import numpy as np
import scipy.io.wavfile as wav
import scipy.signal as sig
import matplotlib.pyplot as plt

# 读取数据文件
fs, signal = wav.read('./train_1s/10/1.wav')

# FIR带通滤波
wn = np.array([10, 1000]) * 2 / fs
b = sig.firwin(101, wn)
signal_f = sig.lfilter(b, 1, signal)

# LOFAR
Tfft = 0.5  # 单个快拍积分时间/秒
nfft = int(Tfft * fs)  # 单个快拍积分采样点数量

L_scan = 0.5 * nfft  # 积分窗滑移长度
N = int(len(signal_f) // L_scan - int(nfft / L_scan))  # 积分滑移次数
df = 1 / Tfft  # 频率分辨率
f = np.arange(10, 1000 + df, df)  # 显示频率范围

p = np.zeros((N, len(f)))
for i in range(1, N):
    data = signal_f[int(1 + (i - 1) * N):int(1 + (i - 1) * N + nfft)]
    f_pxx, pxx = sig.periodogram(data, fs, window='hamming', nfft=nfft)
    # 选择与f相同范围的频率和功率谱
    idx = np.where((f_pxx >= f[0]) & (f_pxx <= f[-1]))[0]
    f_pxx = f_pxx[idx]
    pxx = pxx[idx]
    p[i, :] = 10 * np.log10(pxx)

tt = np.floor((np.arange(N, 0, -1) * L_scan + nfft / 2) / fs)  # 时间序列
plt.figure()
plt.imshow(p, aspect='auto', extent=[f[0], f[-1], tt[0], tt[-1]])
plt.colorbar()
plt.show()